/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"
#include <iostream>

using namespace std;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid = -1;
    treeHeight = 0;
    modified = 0;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
	// check if can open indexname
	if((pf.open(indexname, mode))!=0){
		return RC_FILE_OPEN_FAILED;
	}

	// if write, no checks needed
	if(mode == 'w'){
		return 0;
	}

	// check if can read file
	int buffer[256];
	if(pf.read(0, buffer) != 0){
		return RC_FILE_READ_FAILED;
	}

	// need to get the rootpid and treeheight from the buffer
	if(modified==0){
		rootPid = buffer[4];
		treeHeight = buffer[5];
		nodetype = buffer[6]; // leafnode or nonleaf node
		modified = 1;
	}else{
		nodetype = buffer[6];
	}
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
    int buffer[256];
    if(mode == 'w'){
    	// store the current rootpid
    	buffer[4] = rootPid;
    	// store current treeheight
    	buffer[5] = treeHeight;
    	if( pf.write(0, buffer) != 0) {
			//cout << "Write to root node failed" << endl;
			return RC_FILE_WRITE_FAILED;
		}
    }

    // nothing to take note of if mode was read, 'r'
    return pf.close();
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
	IndexCursor cursor;
	locate(key, cursor); // locates key or closest spot and stores in cursor

	BTLeafNode leaf;

	PageId tempP;
	PageId oPid;
	BTNonLeafNode sibling_nlnode;
	int midKey;
	BTNonLeafNode node;

	if(treeHeight == 1){
		// root is leaf
		leaf.read(cursor.pid, pf);
		if(leaf.insert(key, rid) != 0){
			// use insert and split here
			BTLeafNode sibling;
			int siblingKey;

			// first insert into leaf node
			leaf.insertAndSplit(key, rid, sibling, siblingKey);
			
			//set pointers
			PageId nextNodePtr = leaf.getNextNodePtr();
			sibling.setNextNodePtr(nextNodePtr);
			leaf.setNextNodePtr(pf.endPid());
			
			// write leaf which contains the first half of the split
			// into cursor.pid in PageFile pf
			if(leaf.write(cursor.pid, pf) != 0){
				return RC_FILE_WRITE_FAILED;
			}
			
			// write sibling which contains second half of the split
			// into pf.endPid() in PageFile pf
			// this is because endPid() returns the last pid in pf
			if(sibling.write(pf.endPid(), pf) != 0){
				return RC_FILE_WRITE_FAILED;
			}

			// now we need to create a new root
			// that points to the leaf and sibling nodes
			// with the siblingKey as the only key in between pageIds
			BTNonLeafNode root;
			root.initializeRoot(cursor.pid, siblingKey, pf.endPid()-1);
			if(root.write(pf.endPid(), pf) != 0){
				return RC_FILE_WRITE_FAILED;
			}
			rootPid = pf.endPid() - 1;
			treeHeight++;
			return 0;
		}
		leaf.write(cursor.pid, pf);
		return 0;
	}

	// for any other treeheights
	// insert
	if(leaf.insert(key, rid) != 0){

		BTLeafNode sibling;
		int siblingKey;

		// first insert into leaf node
		leaf.insertAndSplit(key, rid, sibling, siblingKey);
		
		//set pointers to link them
		PageId nextNodePtr = leaf.getNextNodePtr();
		sibling.setNextNodePtr(nextNodePtr);
		leaf.setNextNodePtr(pf.endPid());
		
		// write leaf which contains the first half of the split
		// into cursor.pid in PageFile pf
		if(leaf.write(cursor.pid, pf) != 0){
			return RC_FILE_WRITE_FAILED;
		}
		
		// write sibling which contains second half of the split
		// into pf.endPid() in PageFile pf
		// this is because endPid() returns the last pid in pf
		if(sibling.write(pf.endPid(), pf) != 0){
			return RC_FILE_WRITE_FAILED;
		}

		// then see if it overflows in parent node
		bool first_iter = true;
		do {
			oPid = tempP;
			tempP = pageNode.back();
			pageNode.pop_back();
			
			if(node.read(tempP, pf) != 0){
				return RC_FILE_READ_FAILED;
			}
			
			if(first_iter) {
				// the first iteration
				if(node.insert(siblingKey, pf.endPid()-1) == 0) {
					// insert works, no need for splitting
					if(node.write(tempP, pf) != 0) {
						return RC_FILE_WRITE_FAILED;
					}
					break; // get out of loop
				}else{
					// must insert and split here
					node.insertAndSplit(siblingKey, cursor.pid, sibling_nlnode, midKey);
					sibling_nlnode.write(pf.endPid(), pf);
					if(node.write(tempP, pf) != 0){
						return RC_FILE_WRITE_FAILED;
					}
				}
			}else{
				// not first iteration
				if(node.insert(midKey, oPid) == 0){
					// insert works
					if(node.write(tempP, pf) != 0){
						return RC_FILE_WRITE_FAILED;
					}
					break; // get out of loop
				}else{
					node.insertAndSplit(midKey, oPid, sibling_nlnode, midKey);
					sibling_nlnode.write(pf.endPid(), pf);
					
					// root case splitting
					if(tempP == rootPid){
						BTNonLeafNode root;
						root.initializeRoot(tempP, midKey, pf.endPid()-1);
						if(root.write(pf.endPid(), pf) != 0){
							return RC_FILE_WRITE_FAILED;
						}
						rootPid = pf.endPid() - 1;
						treeHeight++;
					}
					if(node.write(tempP, pf) != 0){
						return RC_FILE_WRITE_FAILED;
					}
					break; // get out of loop
				}
			}
			first_iter = false;
		}while(tempP != rootPid);

	}

	// for any other treeheights without need of 
	// insert and split, go here
	if(leaf.write(cursor.pid, pf) != 0){
		// write failed
		return RC_FILE_WRITE_FAILED;
	}
    return 0;
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
	
	if(treeHeight == 0){
		// tree height == 0 means nothing is inside the index
		cursor.pid = 1;
		cursor.eid = 0;
		treeHeight++;
		rootPid = 1;
	}else if(treeHeight == 1){
		//tree height == 1 means only leaf type root exists
		BTLeafNode root;
		int eid;
		if(root.read(rootPid, pf) != 0){
			return RC_FILE_READ_FAILED;
		}
		
		if(root.locate(searchKey, eid) != 0){
			return RC_NO_SUCH_RECORD;
		}
		cursor.pid = rootPid;
		cursor.eid = eid;
	}else{
		//non-leaf nodes exist
		BTNonLeafNode node;
		if(node.read(rootPid, pf) != 0){
			return RC_FILE_READ_FAILED;
		}
		
		int height = 1;
		while(height < treeHeight){
			//we insert the nextPid to our vector to maintain it.
			if(height == 1) {
				pageNode.push_back(rootPid);
			}else{
				pageNode.push_back(nextPid);
			}
			
			if(node.locateChildPtr(searchKey, nextPid) != 0){
				return RC_NO_SUCH_RECORD;
			}
			
			
			if(node.read(nextPid, pf) != 0){
				return RC_FILE_READ_FAILED;
			}
			height++;
		}
		
		// still need to go through the leaf node entries
		cursor.pid = nextPid;
		int eid;
		BTLeafNode leafy;
		leafy.read(nextPid, pf);
		leafy.locate(searchKey, eid);
		cursor.eid = eid;
	}
	
    return 0;
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
	BTLeafNode node;
	// read pid from pf 
	if(node.read(cursor.pid, pf) != 0){
		return RC_FILE_READ_FAILED;
	}
	
	// read entry from the pid
	if(node.readEntry(cursor.eid, key, rid) != 0){
		return RC_INVALID_RID;
	}
	
	// increment cursor or move to next pid
	if(cursor.eid < node.getKeyCount() - 1){
		cursor.eid++;
	}else{
		cursor.eid = 0;
		if(node.gethasnext() == 0){ 
			return RC_END_OF_TREE;
		}
		cursor.pid = node.getNextNodePtr();
	}	
	
    return 0;    
}
